{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b99a6e75-2593-4def-b695-b7a031a392a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue May  7 09:14:18 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 536.67                 Driver Version: 536.67       CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                     TCC/WDDM  | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA GeForce RTX 3080 Ti   WDDM  | 00000000:01:00.0  On |                  N/A |\n",
      "|  0%   46C    P8              43W / 350W |    591MiB / 12288MiB |      6%      Default |\n",
      "|                                         |                      |                  N/A |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|    0   N/A  N/A      1500    C+G   C:\\Windows\\System32\\dwm.exe               N/A      |\n",
      "|    0   N/A  N/A      6240    C+G   C:\\Windows\\explorer.exe                   N/A      |\n",
      "|    0   N/A  N/A      8116    C+G   ...5n1h2txyewy\\ShellExperienceHost.exe    N/A      |\n",
      "|    0   N/A  N/A     10180    C+G   ...oogle\\Chrome\\Application\\chrome.exe    N/A      |\n",
      "|    0   N/A  N/A     11692    C+G   ....Search_cw5n1h2txyewy\\SearchApp.exe    N/A      |\n",
      "|    0   N/A  N/A     12268    C+G   ...h2txyewy\\InputApp\\TextInputHost.exe    N/A      |\n",
      "|    0   N/A  N/A     13308    C+G   ...GeForce Experience\\NVIDIA Share.exe    N/A      |\n",
      "|    0   N/A  N/A     16012    C+G   ...9\\extracted\\runtime\\WeChatAppEx.exe    N/A      |\n",
      "|    0   N/A  N/A     16168    C+G   ...rm 2022.1.2\\jbr\\bin\\jcef_helper.exe    N/A      |\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "067cec4c-2c3d-4d0a-9999-2a9ada4bd52b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cpu'), device(type='cuda'), device(type='cuda', index=1))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "torch.device(\"cpu\"),torch.device(\"cuda\"),torch.device(\"cuda:1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "56578403-9340-4fce-b496-0740ccaa8b18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fe295769-16ba-49ec-a4dc-bc4f12d4fd5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cuda', index=0),\n",
       " device(type='cpu'),\n",
       " [device(type='cuda', index=0)])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def try_gpu(i=0):\n",
    "    if torch.cuda.device_count() >= i+1:\n",
    "        return torch.device(f\"cuda:{i}\")\n",
    "    return torch.device(\"cpu\")\n",
    "\n",
    "def try_all_gpus():\n",
    "    devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]\n",
    "    return devices if devices else [torch.device(\"cpu\")]\n",
    "try_gpu(),try_gpu(10),try_all_gpus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7a83b831-f2b8-4999-822e-437aa6a73616",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#5.6.2. 张量与GPU\n",
    "x = torch.tensor([1,2,3])\n",
    "x.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "07ad7368-71c2-45d6-89c8-cf902b341b52",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1.],\n",
       "        [1., 1., 1.]], device='cuda:0')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#5.6.2.1. 存储在GPU上\n",
    "x = torch.ones(2,3,device=try_gpu())\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6690dec7-3391-42f8-9f7e-c64c80873b87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.]], device='cuda:0')\n",
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "#5.6.2.2. 复制\n",
    "z = x.cuda(0)\n",
    "print(x)\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b9479dc9-0e6c-4ff9-b65b-e0ae9db20c39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2., 2., 2.],\n",
       "        [2., 2., 2.]], device='cuda:0')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3ee7ead5-2479-4f86-b1a0-643351a349b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.cuda(0) is z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "adec710b-2d2e-4482-a079-9a0814b55503",
   "metadata": {},
   "outputs": [],
   "source": [
    "#5.6.3. 神经网络与GPU\n",
    "net = nn.Sequential(nn.Linear(3,1))\n",
    "net = net.to(device=try_gpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5707f646-7d3d-4488-b85f-44135d0f136c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.6327],\n",
       "        [0.6327]], device='cuda:0', grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b08232ca-b347-49bb-94a7-66787363a519",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net[0].weight.data.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c0ea28ea-634e-4d53-a7df-e4ff14efc703",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu计算总时长: 151.0 ms\n",
      "gpu计算总时长: 4.03 ms\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import torch\n",
    "\n",
    "def try_gpu(i=0):  #@save\n",
    "    \"\"\"如果存在，则返回gpu(i)，否则返回cpu()\"\"\"\n",
    "    if torch.cuda.device_count() >= i + 1:\n",
    "        return torch.device(f'cuda:{i}')\n",
    "    return torch.device('cpu')\n",
    "\n",
    "startTime1=time.time()\n",
    "for i in range(100):\n",
    "    A = torch.ones(500,500)\n",
    "    B = torch.ones(500,500)\n",
    "    C = torch.matmul(A,B)\n",
    "endTime1=time.time()\n",
    "\n",
    "startTime2=time.time()\n",
    "for i in range(100):\n",
    "    A = torch.ones(500,500,device=try_gpu())\n",
    "    B = torch.ones(500,500,device=try_gpu())\n",
    "    C = torch.matmul(A,B)\n",
    "endTime2=time.time()\n",
    "\n",
    "print('cpu计算总时长:', round((endTime1 - startTime1)*1000, 2),'ms')\n",
    "print('gpu计算总时长:', round((endTime2 - startTime2)*1000, 2),'ms')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1c3d6ce-95e9-4c59-b468-eb2dfa039aaa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
